import torch.nn as nn
from TorchCRF import CRF


class NERBiLSTM_CRF(nn.Module):
    def __init__(self, embedding_dim, hidden_dim, dropout, word2id, tag2id):
        super(NERBiLSTM_CRF, self).__init__()
        self.name = "BiLSTM_CRF"
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.vocab_size = len(word2id) + 1
        self.tag_size = len(tag2id)
        self.tag_to_ix = tag2id
        # 定义lstm网络层
        self.embed = nn.Embedding(self.vocab_size, self.embedding_dim)
        self.dropout = nn.Dropout(dropout)

        self.lstm = nn.LSTM(self.embedding_dim, self.hidden_dim // 2,
                            bidirectional=True, batch_first=True)

        self.out = nn.Linear(self.hidden_dim, self.tag_size)
        # 定义CRF网络层
        self.crf = CRF(self.tag_size)

    def forward(self, x, mask):
        """
        预测最终结果，不参与训练，预测时使用。
        :param x: 输入数据。
        :param mask: 掩码。
        :return: 预测结果。
        """
        # lstm模型得到的发射分数
        outputs = self.get_lstm2linear(x)
        # 只保留有效位置的发射分数
        outputs = outputs * mask.unsqueeze(-1)
        # 发射分数送入维特比解码
        outputs = self.crf.viterbi_decode(outputs, mask)
        return outputs

    def log_likelihood(self, x, tags, mask):
        """
        计算损失结果，训练时使用。
        损失函数：CRF损失——真实路径分数/所有路径分数
        :param x: 输入数据
        :param tags: 真实标签
        :param mask: 掩码，可以不加。
        :return: 损失结果
        """
        # lstm模型得到的结果
        outputs = self.get_lstm2linear(x)
        outputs = outputs * mask.unsqueeze(-1)
        # 计算损失
        return - self.crf(outputs, tags, mask)

    def get_lstm2linear(self, x):
        embed_x = self.embed(x)
        outputs, hidden = self.lstm(embed_x)
        outputs = self.dropout(outputs)
        outputs = self.out(outputs)
        return outputs
